﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using AzureMRCore.DataCommunication;
using AzureMRCore.DataModel;
using AzureMRCore.DataTypes;
using AzureMRCore.Drivers;
using AzureMRCore.InputFormat;
using AzureMRCore.IterativeDataCache;
using AzureMRCore.OutputCollectors;
using AzureMRCore.Partitioners;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;

namespace AzureMRCore.MapRed
{
    public class Mapper<INKEY, INVALUE, OUTKEY, OUTVALUE, BCASTINKEY, BCASTINVALUE> : IMapper
        where INKEY : Key, new()
        where OUTKEY : Key, new()
        where INVALUE : IValue, new()
        where OUTVALUE : IValue, new()
        where BCASTINKEY : Key, new()
        where BCASTINVALUE : IValue, new()
    {
        private DataCache _dataCache;
        private MRConf _mrConf;
        private MRDriver _mrDriver;
        private CloudStorageAccount _storageAccount;
        private CloudBlobClient _blobClient;
        protected Guid _workerId;
        protected string Stats = "";

        public Boolean Init(CloudStorageAccount storageAccount, MRConf mrConf, MRDriver mrDriver, DataCache dataCache,
                            object configOut)
        {
            _mrConf = mrConf;
            _mrDriver = mrDriver;
            _storageAccount = storageAccount;
            _dataCache = dataCache;
            _blobClient = _storageAccount.CreateCloudBlobClient();

            _workerId = Guid.NewGuid();
            Config(storageAccount, configOut);
            return true;
        }

        protected virtual int Map(INKEY key, INVALUE value, List<KeyValuePair<BCASTINKEY, BCASTINVALUE>> dynamicData, 
            IOutputCollector<OUTKEY, OUTVALUE> outputCollector, string programArgs)
        {
            return 0;
        }

        protected virtual Boolean Config(CloudStorageAccount storageAccount, object configOut)
        {
            return true;
        }

        /// <summary>
        ///   Scheduled from Cache
        /// </summary>
        /// <param name="inputFormat"> </param>
        /// <param name="mapTaskContext"> </param>
        /// <returns> </returns>
        public bool ExecuteMapTask(ICachedInputFormat inputFormat, MapTaskContext mapTaskContext)
        {
            return ExecuteMapTask((IInputFormat<INKEY, INVALUE>) inputFormat, mapTaskContext);
        }

        /// <summary>
        ///   scheduled from queue
        /// </summary>
        /// <param name="mapTaskContext"> </param>
        /// <returns> </returns>
        public bool ExecuteMapTask(MapTaskContext mapTaskContext)
        {
            IInputFormat<INKEY, INVALUE> inputFormat = null;
            if (mapTaskContext.IsBlobType)
            {
                inputFormat = _dataCache.GetData(mapTaskContext.BlobUrl,
                                                 _mrDriver.GetInputFormat<INKEY, INVALUE>(),
                                                 mapTaskContext.JobID, mapTaskContext.Iteration,
                                                 mapTaskContext.AppName,
                                                 _workerId.ToString(), mapTaskContext.MapID);
            }
            return ExecuteMapTask(inputFormat, mapTaskContext);
        }

        private bool ExecuteMapTask(IInputFormat<INKEY, INVALUE> inputFormat, MapTaskContext mapTaskContext)
        {
            int status = -1;
            List<KeyValuePair<BCASTINKEY, BCASTINVALUE>> bCastData = null;

            MapDataModel mapTask = mapTaskContext.MapTask;

            Stopwatch executionTimer = Stopwatch.StartNew();
            bCastData = GetBCastData(mapTaskContext, mapTask, bCastData);
            mapTask.BroadcastFetchTime = executionTimer.ElapsedMilliseconds;

            var combiner = _mrDriver.GetCombiner();
            var partitioner = _mrDriver.PartitionerType;
            partitioner.NumReducers = mapTaskContext.NumReduceTasks;
            var outputCollector = combiner == null
                                      ? (IOutputCollector<OUTKEY, OUTVALUE>)
                                        new MapOutputCollector<OUTKEY, OUTVALUE>(partitioner)
                                      : new MapCombineOutputCollector<OUTKEY, OUTVALUE>();

            if (!mapTaskContext.IsCustomType)
            {
                //execute the map.. Output should be the file name and the file URL or key and value strings  
                executionTimer.Restart();
                while (inputFormat.HasNext())
                {
                    if (inputFormat.Next())
                    {
                        INVALUE value = inputFormat.GetCurrentValue();
                        INKEY key = inputFormat.GetCurrentKey();
                        status = Map(key, value, bCastData, outputCollector, mapTaskContext.ProgramParams);
                    }
                }
                inputFormat.Close();
                mapTask.ExecTime = executionTimer.ElapsedMilliseconds;
                executionTimer.Restart();
            }
            else
            {
                executionTimer = Stopwatch.StartNew();

                INKEY inKey = new INKEY();
                inKey.Parse(mapTaskContext.CustomKey);
                INVALUE inValue = new INVALUE();
                inValue.ParseValue(mapTaskContext.CustomValue);

                status = Map(inKey, inValue, bCastData, outputCollector, mapTaskContext.ProgramParams);
                executionTimer.Stop();
            }

            if (combiner != null)
            {
                outputCollector = CombineMapOutput(combiner, partitioner, mapTaskContext.ProgramParams, outputCollector,
                                                   ref status);
                mapTask.CombinerTime = executionTimer.ElapsedMilliseconds;
            }

            executionTimer.Restart();

            if (status == -1)
                return false;

            byte[] reduceProductCounts = ShuffleMapOutputData(outputCollector, mapTaskContext,
                                                              _mrDriver);

            //update the map task status in the table  TODO : can we do all these at the end                         
            mapTaskContext.UpdateMapTaskStatus(MapDataModel.FINISHED, reduceProductCounts);
            mapTask.OutputTime = executionTimer.ElapsedMilliseconds;
            mapTask.Attempt3StartTime = Stats;
            Stats = "";
            return true;
        }

        public byte[] ShuffleMapOutputData(IOutputCollector<OUTKEY, OUTVALUE> outputCollector,
                                           MapTaskContext mapTaskContext, MRDriver mrDriver)
        {
            var outFormat = mrDriver.GetMapOutputFormat<OUTKEY, OUTVALUE>();
            var mapOutCollector = (MapOutputCollector<OUTKEY, OUTVALUE>) outputCollector;
            // To keep track of the total number of reduce data products.
            byte[] reduceProductCounts = new byte[0];
            var dataCommunicator = DataCommunicator.GetInstance(_mrConf, _storageAccount);

            /*
             * handle primitive
             * handle single result
             * handle multi results
             */
            PrimitiveConf primitiveConf = mrDriver.GetPrimitive(Int32.Parse(mapTaskContext.Iteration));
            if (primitiveConf.IsPrimitive)
            {
                dataCommunicator.ExecuteAllGatherTransfer<OUTKEY, OUTVALUE, BCASTINKEY, BCASTINVALUE>(mapTaskContext,
                                                                                                      mrDriver,
                                                                                                      primitiveConf,
                                                                                                      mapOutCollector);
            }
            if (mapTaskContext.NumReduceTasks == 0)
            {
                foreach (KeyValuePair<int, List<KeyValuePair<OUTKEY, OUTVALUE>>> record in mapOutCollector.GetResults())
                {
                    string containerName = mapTaskContext.OutDir.ToLower();
                    CloudBlobContainer container = _blobClient.GetContainerReference(containerName);
                    container.CreateIfNotExist();
                    string firstID = mapTaskContext.MapTask.RowKey + "_" + record.Key;
                    outFormat.UploadValues(record.Value, firstID, mapTaskContext.MapTask.PartitionKey, container);
                }
            }
            else
            {
                reduceProductCounts = new byte[mapTaskContext.NumReduceTasks];
                foreach (var record in mapOutCollector.GetResults())
                {
                    List<KeyValuePair<OUTKEY, OUTVALUE>> kvPairsForReducers = record.Value;
                    // if it's a single value and less than 64kb, then send in the table
                    if (kvPairsForReducers.Count == 1 &&
                        kvPairsForReducers.First().Value.Size() < (63 * 1024))
                    {
                        byte[] data = kvPairsForReducers.First().Value.GetBytes();
                        reduceProductCounts = dataCommunicator.TransmitSingleValue(mapTaskContext, data, record.Key,
                                                                                   kvPairsForReducers.First().Key.GetTextValue());
                    }
                    else
                    {
                        dataCommunicator.TransmitMultipleValues(mapTaskContext, outFormat, _mrDriver.DirectTCPTranfer,
                                                                reduceProductCounts, record);
                    }
                }
            }
            return reduceProductCounts;
        }

        /// <summary>
        ///   Local combiner step for map tasks
        /// </summary>
        /// <param name="combinerType"> </param>
        /// <param name="partitioner"> </param>
        /// <param name="programParams"> </param>
        /// <param name="outputCollector"> </param>
        /// <param name="status"> </param>
        /// <returns> </returns>
        private static IOutputCollector<OUTKEY, OUTVALUE> CombineMapOutput(IReducer combinerType,
                                                                           IPartitioner partitioner,
                                                                           string programParams,
                                                                           IOutputCollector<OUTKEY, OUTVALUE>
                                                                               outputCollector,
                                                                           ref int status)
        {
            var mapCombineOutputCollector = (MapCombineOutputCollector<OUTKEY, OUTVALUE>) outputCollector;
            Dictionary<OUTKEY, List<OUTVALUE>> mapResults = mapCombineOutputCollector.GetResults();

            var combinerOutCollector = new MapOutputCollector<OUTKEY, OUTVALUE>(partitioner);
            var combiner = (Reducer<OUTKEY, OUTVALUE, OUTKEY, OUTVALUE>) combinerType;

            foreach (KeyValuePair<OUTKEY, List<OUTVALUE>> recordSet in mapResults)
            {
                status = combiner.Reduce(recordSet.Key, recordSet.Value, combinerOutCollector, programParams);
            }
            outputCollector = combinerOutCollector;
            return outputCollector;
        }

        private List<KeyValuePair<BCASTINKEY, BCASTINVALUE>> GetBCastData(MapTaskContext mapTaskContext,
                                                                          MapDataModel mapTask,
                                                                          List<KeyValuePair<BCASTINKEY, BCASTINVALUE>>
                                                                              bCastData)
        {
            if (mapTaskContext.IsPrimitiveBCast)
            {
                var bCastReceiver = _mrConf.BCastReceiver;
                //wait till the deamon fetch the data
                while (!bCastReceiver.ContainsBCastData(mapTask.BroadcastData))
                {
                    Thread.Sleep(40);
                }
                var bCastInputFormat =
                    (IInputFormat<BCASTINKEY, BCASTINVALUE>) bCastReceiver.GetBCastData(mapTask.BroadcastData);
                bCastData = bCastInputFormat.GetAllPairs();
            }
                // We can get rid of the following when we move all BCast to receiver daemon
            else if (mapTask.BroadcastData != null)
            {
                IInputFormat<BCASTINKEY, BCASTINVALUE> bCastInputFormat =
                    _dataCache.GetData(mapTask.BroadcastData,
                                       _mrDriver.GetBCastInputFormat<BCASTINKEY, BCASTINVALUE>(),
                                       mapTaskContext.JobID, mapTaskContext.Iteration, mapTask.AppName,
                                       _workerId.ToString(), null, true);
                bCastData = bCastInputFormat.GetAllPairs();
                Task cleanupTask = new Task(() => _dataCache.CleanupBCastData(mapTaskContext.JobID, mapTask.AppName,
                                                                              Int32.Parse(mapTaskContext.Iteration)));
                cleanupTask.Start();
            }
            return bCastData;
        }
    }
}